import tifffile
import numpy as np
from Data.DataStorage import HDFS
import cv2
from PIL import Image, ImageOps
Image.MAX_IMAGE_PIXELS = None
import os
from Util.Public import *
from Util.Process import INFO


def ensure_local_data(root_path, tiff_to_other_data_flag=False):
    hdfs_path = "/" + root_path
    local_path = get_path(os.path.abspath(__file__), 2, root_path)
    if not os.path.exists(local_path):
        local_dir = os.path.dirname(local_path)
        hdfs = HDFS()
        hdfs.pull(hdfs_path, local_dir)
    if tiff_to_other_data_flag:
        file_name_with_ext = os.path.basename(local_path)
        file_name_not_ext = os.path.splitext(file_name_with_ext)[0]
        local_jpg_path = get_path(os.path.abspath(__file__), 2, get_config("data", "TASK")['raw_path'] + "/" + file_name_not_ext + ".jpg")
        tiff_to_other_data(local_path, local_jpg_path, True)
        return local_jpg_path
    else:
        return local_path

def resize_image(image_path):
    image = Image.open(image_path)
    width, height = image.size
    new_width = width // 2
    new_height = height // 2
    if new_width > 10000 or new_height > 10000:
        resized_image = image.resize((new_width, new_height))
        resized_image.save(image_path)


def tiff_to_other_data(tiff_path, other_data_path, resize_flag=False):
    dataset = get_dataset(os.path.basename(tiff_path))
    if dataset == "Landsat":
        mode = 0
    elif dataset == "Sentinel":
        mode = 1
    else:
        return
    if os.path.exists(other_data_path):  # 如果存在则先删除
        os.remove(other_data_path)

    if mode == 0:  # For Landsat
        png_data_path = os.path.splitext(other_data_path)[0] + '.png'
        original_image = cv2.imread(tiff_path)
        # 裁剪黑色边界
        gray = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)
        # 查找图像边界
        _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # 寻找最大的边界框
        max_contour = max(contours, key=cv2.contourArea)
        x, y, w, h = cv2.boundingRect(max_contour)
        # 裁剪图像
        cropped_image = original_image[y:y + h, x:x + w]
        # 转换为灰度图像
        gray_image = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2GRAY)
        # 直方图均衡化
        equalized_image = cv2.equalizeHist(gray_image)
        rgb_image = Image.fromarray(cv2.cvtColor(equalized_image, cv2.COLOR_BGR2RGB))
        # 保存为PNG格式
        rgb_image.save(png_data_path, format="PNG")
        os.rename(png_data_path, other_data_path)

    if mode == 1:  # For Sentinel
        data = tifffile.imread(tiff_path)
        data = cv2.flip(data, 1)
        data = cv2.rotate(data, cv2.ROTATE_180)
        height, width = data.shape[:2]
        if resize_flag:
            data = cv2.resize(data, (width // 2, height // 2))
        cv2.imwrite(other_data_path, data)
    INFO("INFO", f"Tiff_to_jpg from {tiff_path} to {other_data_path}")


def split_tile(input_path, output_path):
    """
    将TIFF图像文件切割成1024x1024分辨率的小图像,并按行列编号命名。

    参数:
    input_path (str): 原始TIFF图像文件的输入路径
    output_path (str): 输出小图像的目录路径
    """

    # 创建输出目录(如果不存在)
    os.makedirs(output_path, exist_ok=True)

    # 打开TIFF图像文件
    with Image.open(input_path) as img:
        width, height = img.size

        # 计算行数和列数
        num_rows = (height + 1023) // 1024
        num_cols = (width + 1023) // 1024

        # 遍历每一行和每一列
        for row in range(num_rows):
            for col in range(num_cols):
                # 计算当前小图像的左上角和右下角坐标
                left = col * 1024
                top = row * 1024
                right = min(left + 1024, width)
                bottom = min(top + 1024, height)

                # 裁剪出当前小图像
                tile = img.crop((left, top, right, bottom))

                # 如果当前小图像的尺寸小于1024x1024，创建一个新的黑色图像并将小图像粘贴在左上角
                if tile.width < 1024 or tile.height < 1024:
                    pad_width = max(0, 1024 - tile.width)
                    pad_height = max(0, 1024 - tile.height)
                    # 在左上角用黑色填充小图像
                    tile = ImageOps.expand(tile, (0, 0, pad_width, pad_height), fill='black')

                # 构造输出文件名
                output_filename = f"{row}_{col}.TIFF"
                output_file = os.path.join(output_path, output_filename)

                # 保存当前小图像
                tile.save(output_file)

        INFO("INFO", f"Split_tiff from {input_path} to {output_path}")


def merge_tile(tile_folder, output_path, tile_size=1024):
    """
    将指定文件夹中的小图像按照行列顺序合并成大图像。

    参数:
    tile_folder (str): 包含小图像的文件夹路径
    output_path (str): 合并后大图像的输出路径
    tile_size (int): 每个小图像的尺寸（默认为1024）
    """

    # 获取文件夹中的所有图像文件
    tile_files = [f for f in os.listdir(tile_folder) if
                  f.endswith('.jpg') or f.endswith('.png') or f.endswith('.tif') or f.endswith('.tiff') or f.endswith(
                      '.TIF') or f.endswith('.TIFF')]

    # 按文件名排序图像文件，确保它们按照行列顺序排列
    tile_files.sort(key=lambda x: tuple(map(int, x.split('.')[0].split('_'))))
    if len(tile_files) == 0:
        INFO("ERROR", f"Merge_tiff: Don't have tiff!")
        return
    # 初始化行和列的最小值和最大值
    min_row = float('inf')
    min_col = float('inf')
    max_row = -float('inf')
    max_col = -float('inf')

    # 遍历列表并提取行和列信息
    for tile in tile_files:
        row, col = map(int, tile.split('.')[0].split('_'))
        min_row = min(min_row, row)
        min_col = min(min_col, col)
        max_row = max(max_row, row)
        max_col = max(max_col, col)

    # 计算大图像的最终尺寸
    final_width = tile_size * (max_col - min_col + 1)
    final_height = tile_size * (max_row - min_row + 1)

    # 创建一个空白的大图像来存储合并后的结果
    result = Image.new('RGB', (final_width, final_height), (0, 0, 0))

    # 遍历每个小图像并合并到大图像中
    for tile_index, tile_filename in enumerate(tile_files):
        # 计算当前小图像所在的行列索引
        row, col = map(int, tile_files[tile_index].split('.')[0].split('_'))
        row = row - min_row
        col = col - min_col
        # 构造当前小图像的路径
        tile_path = os.path.join(tile_folder, tile_filename)
        tile_img = Image.open(tile_path)

        # 将小图像粘贴到大图像的对应位置
        left = col * tile_size
        top = row * tile_size

        result.paste(tile_img, (left, top))

    # 保存合并后的大图像
    result.save(output_path)
    INFO("INFO", f"Saved merged image from {tile_folder} to {output_path}")


def filter_tile(folder_path, range_str):
    # 解析行列范围字符串
    parts = range_str.split("_")
    min_row, max_row = map(int, parts[:2])
    min_col, max_col = map(int, parts[2:])

    # 遍历文件夹中的文件
    for filename in os.listdir(folder_path):
        # 解析文件名中的行列信息
        parts = filename.split("_")
        row = int(parts[0])
        col = int(parts[1].split(".")[0])  # 移除文件扩展名后的部分

        # 判断是否在指定范围内，不在范围内则删除文件
        if not (min_row <= row <= max_row and min_col <= col <= max_col):
            file_path = os.path.join(folder_path, filename)
            os.remove(file_path)
